import math
import numpy as np
from scipy import signal
from matplotlib import pyplot as plt

"""
{
    "peak": float,  # 峰值
    "peak_peak": float,  # 峰峰值
    "mean": float,  # 均值
    "rms": float,  # 均方根值（有效值）
    "std": float,  # 标准差
    "var": float,  # 方差
    "skewness": float,  # 偏度
    "kurtosis": float,  # 峭度
    "waveform_factor": float,  # 波形指标
    "crest_indicator": float,  # 峰值指标
    "impulse_indicator": float,  # 脉冲指标
    "clearance_indicator": float,  # 裕度指标
    "fc": float,  # 重心频率
    "msf": float,  # 均方频率
    "vf": float,  # 频率方差
    "speed_rms": float,  # 速度有效值
    "GB_a_threshold": float,  # 国标规定的加速度阈值
    "GB_v_threshold": float  # 国标规定的速度阈值
    "fault_category": int  # 故障类别, 0-内圈, 1-外圈, 2-滚动体, 3-正常
}
"""

def process_acceleration_data(data: np.ndarray, fs: int) -> tuple:
    """
    处理加速度数据，提取时域、频域和速度有效值等特征

    :param data: 加速度数据
    :param fs: 采样频率
    :return: 特征值
    """
    # 电压转换为加速度，单位是m/s^2
    data = data / 10.3

    # 时域特征

    n = len(data)
    # 1.计算峰值
    peak = max(data)
    # 2.计算峰峰值
    peak_peak = peak - min(data)
    # 3.计算均值
    mean = sum(data) / n
    # 4.计算均方根值（有效值）
    rms = math.sqrt(sum([x ** 2 for x in data]) / n)
    # 5.计算标准差
    std = math.sqrt(sum([(x - mean) ** 2 for x in data]) / n)
    # 6.计算方差
    var = std ** 2
    # 7.计算偏度
    skewness = (sum([(x - mean) ** 3 for x in data]) / n) / (std ** 3)
    # 8.计算峭度, 正态分布的峭度是3, 减去3后得到的结果是相比于正态分布的峭度偏离值
    kurtosis = (sum([(x - mean) ** 4 for x in data]) / n) / (std ** 4) - 3
    abs_data = [abs(x) for x in data]
    # 9.计算波形指标
    waveform_factor = rms / (sum(abs_data) / n)
    # 10.计算峰值指标
    crest_indicator = max(abs_data) / rms
    # 11.计算脉冲指标
    impulse_indicator = max(abs_data) / (sum(abs_data) / n)
    # 12.计算裕度指标
    clearance_indicator = max(abs_data) / (sum([x ** 0.5 for x in abs_data]) / n) ** 2

    # 频域特征

    # 傅里叶变换
    ffted = abs(np.fft.fft(data))[:n // 2]
    freq = fs / n * np.arange(n // 2)
    
    # 1.重心频率
    fc = sum([freq[i] * ffted[i] for i in range(n // 2)]) / sum(ffted)
    # 2.均方频率
    msf = sum([freq[i] ** 2 * ffted[i] for i in range(n // 2)]) / sum(ffted)
    # 3.频率方差
    # vf = sum([(freq[i] * ffted[i] - ffted.mean()) ** 2 for i in range(n // 2)]) / sum(ffted)

    # 速度积分
    data_rm_mean = data - mean
    data_rm_mean = np.insert(data_rm_mean, 0, 0)
    dt = 1 / fs
    speed = np.array([(data_rm_mean[i] + data_rm_mean[i + 1]) * dt / 2 for i in range(n)])

    # 位移积分
    speed_rm_mean = speed - speed.mean()
    speed_rm_mean = np.insert(speed_rm_mean, 0, 0)
    displacement = [(data_rm_mean[i] + data_rm_mean[i + 1]) * dt / 2 for i in range(n)]
    vf = math.sqrt(sum([x ** 2 for x in displacement]) / n)
    # 1.速度有效值
    speed_rms = np.power(speed, 2).mean() ** 0.5
    # speed_rms = math.sqrt(sum([x ** 2 for x in speed]) / n)

    # 10Hz低通滤波
    ffted = np.fft.fft(data)
    freq = np.fft.fftfreq(n, 1/fs)
    mask = np.abs(freq <= 10)# & np.abs(freq <= 1000)
    data_10Hz = np.real(np.fft.ifft(ffted * mask))
    data_10Hz = np.insert(data_10Hz, 0, 0)
    GB_a_threshold = (sum([((data_10Hz[i] ** 2 + data_10Hz[i + 1]) ** 2) * dt / 2 for i in range(n)]) / (n / fs)) ** 0.5

    # 10-1000Hz带通滤波
    ffted = np.fft.fft(speed)
    freq = np.fft.fftfreq(n, 1/fs)
    mask = np.abs(freq <= 10) & np.abs(freq <= 1000)
    speed_10_1000Hz = np.real(np.fft.ifft(ffted * mask))
    speed_10_1000Hz = np.insert(speed_10_1000Hz, 0, 0)
    GB_v_threshold = (sum([((speed_10_1000Hz[i] ** 2 + speed_10_1000Hz[i + 1]) ** 2) * dt / 2 for i in range(n)]) / (n / fs)) ** 0.5

    return [
        peak, peak_peak, mean, rms, std, var, skewness, kurtosis, waveform_factor, crest_indicator, impulse_indicator,
        clearance_indicator, fc, msf, vf, speed_rms, GB_a_threshold, GB_v_threshold
    ]



if __name__ == '__main__':
    lst = [0x55, 1]
    print(lst)

